from scapy.all import *
from utils.resolve_packet import *
from utils.save2har import save2json
from utils.setting import get_settings
from threading import Lock
import threading
from common.logger import log
import ast
import dpkt
import jinja2
from scapy_http.http import HTTPRequest, HTTPResponse
from scapy.utils import PcapWriter


def get_host_ip():
    """
    查询本机ip地址
    :return: ip
    """
    import socket
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
    finally:
        s.close()
    return ip


class Capture:
    def __init__(self, config=None):
        self.data = {}
        self.env = {}
        if not config:
            config = 'config.yml'
        self.lock = Lock()
        _setting = get_settings(config)
        self.port = _setting.get('base_filter', {}).get('port')
        self.host = _setting.get('base_filter', {}).get('host')
        self.iface = _setting.get('base_filter', {}).get('iface')
        self._app_filter = _setting.get('app_filter')
        self.init_data()

    def run(self, network_name=None, callback=None):
        if not network_name:
            network_name = self.iface
        if not callback:
            callback = self.save2file
        # _filter = f'host {self.host} and tcp'
        _filter = f"ip src or dst {self.host} and tcp and tcp port {self.port}"

        t1 = threading.Thread(target=sniff,
                              kwargs={'filter': _filter, 'iface': network_name, 'prn': self.pack_callback, 'count': 0},
                              name='sniff')
        t2 = threading.Thread(target=self.read_data, name='read_data', kwargs={'callback': callback})
        t1.start()
        t2.start()
        t1.join()
        t2.join()

    def _duplicate(self, obj: object):
        """
        将对象中包含中划线“-”的属性名称中的中划线修改为下划线“_”,复制新增属性
        解决中划线无法使用obj.xxx-xxx这种调用方式的问题，现在可通过obj.xxx_xxx调用属性
        :param obj:
        :return:
        """
        attr_list = dir(obj)
        for attr in attr_list:
            if "-" in attr:
                new_attr = attr.replace('-', "_")
                setattr(obj, new_attr, getattr(obj, attr))
        return obj

    def filter(self, request, response):
        """
        应用规则过滤
        :return:
        """
        self.env['request'] = self._duplicate(request)
        self.env['response'] = self._duplicate(response)
        faild_cnt = 0
        for engine, rules in self._app_filter.items():
            if not rules:
                continue
            for expect, actual in rules.items():
                cmd = f"self.{engine}('{expect}', '{actual}')"
                try:
                    # log.i('eval({})'.format(cmd))
                    eval(cmd)
                except AssertionError:
                    faild_cnt += 1
        if faild_cnt > 0:
            return False
        else:
            return True

    def _template_render(self, tp_str: str):
        tp_str = '{{ %s }}' % tp_str
        tp = jinja2.Template(tp_str)
        r = tp.render(**self.env)
        if r:
            try:
                r = ast.literal_eval(r)
            except Exception:
                raise
        if isinstance(r, bytes):
            r = r.decode()
        return r

    def eq(self, expect: str, actual: str):
        r = self._template_render(expect)
        log.d('拦截规则： eq, expect: {}, actual: {}; 拦截结果： {}'.format(r, actual, r == actual))
        if r:
            assert r == actual
        else:
            raise AssertionError

    def contain(self, expect: str, actual: str):
        r = self._template_render(expect)
        log.d('拦截规则： contain, expect: {}, actual: {}; 拦截结果： {}'.format(r, actual,
                                                                       actual in r if r is not None else False))
        if r:
            assert actual in r
        else:
            raise AssertionError

    def not_contain(self, expect: str, actual: str):
        r = self._template_render(expect)
        log.d('拦截规则： not_contain, expect: {}, actual: {}; 拦截结果： {}'.format(r, actual,
                                                                           actual not in r if r is not None else False))
        if r:
            assert actual not in r
        else:
            raise AssertionError

    def neq(self, expect: str, actual: str):
        r = self._template_render(expect)
        log.d('拦截规则： neq, expect: {}, actual: {}; 拦截结果： {}'.format(r, actual, actual != r))
        if r:
            assert actual != r
        else:
            raise AssertionError

    def gt(self, expect: str, actual: str):
        r = self._template_render(expect)
        log.d('拦截规则： gt, expect: {}, actual: {}; 拦截结果： {}'.format(r, actual,
                                                                  int(r) > int(actual) if r is not None else False))
        if r:
            assert int(r) > int(actual)
        else:
            raise AssertionError

    def lt(self, expect: str, actual: str):
        r = self._template_render(expect)
        log.d('拦截规则： lt, expect: {}, actual: {}; 拦截结果： {}'.format(r, actual,
                                                                  int(r) < int(actual) if r is not None else False))
        if r:
            assert int(r) < int(actual)
        else:
            raise AssertionError

    def init_data(self):
        self.flag_request = False
        self.flag_response = False
        self.request_raw_data = b''
        self.response_raw_data = b''

        self._response = None
        self._request = None

    def verify_data(self, type, data):
        if type == 'request':
            try:
                _ = dpkt.http.Request(data)
                return True
            except dpkt.dpkt.NeedData:
                return False
        if type == 'response':
            try:
                _ = dpkt.http.Response(data)
                return True
            except dpkt.dpkt.NeedData:
                return False

    def read_data(self, callback=None):
        while True:
            time.sleep(1)
            self.lock.acquire()
            for key in list(self.data.keys()):
                i = 0
                value = self.data[key]
                while i < len(value['request']) and i < len(value['response']):
                    request = value['request'][i]
                    response = value['response'][i]
                    if not self.filter(request=request, response=response):
                        log.i('拦截规则生效，放弃本次请求')
                        del value['request'][i]
                        del value['response'][i]
                        if not value['request']:
                            del self.data[key]
                        continue
                    self.resolve_handler = ResolvePacket()
                    self.resolve_handler.resolve_packet({"request": request, 'response': response})
                    self.resolve_handler.show()
                    if callback and callable(callback):
                        callback(self.resolve_handler._dict)
                    del value['request'][i]
                    del value['response'][i]
                    if not value['request']:
                        del self.data[key]
            self.lock.release()


    def save2file(self, data):
        save2json(data)

    def pack_callback(self, packet):
        # print(type(packet))
        if packet.haslayer("HTTP"):
            IP = packet.payload
            TCP = packet.payload.payload
            HTTP = packet.payload.payload.payload.payload
            if isinstance(HTTP, HTTPRequest):
                log.d('host: {}'.format(IP.dst))
                log.d('method: {}, path: {}'.format(HTTP.Method.decode(), HTTP.Path.decode()))
                session_key = "{}-{}-{}-{}".format(IP.fields.get("src"), IP.fields.get('dst'), TCP.fields.get("sport"),
                                                   TCP.fields.get("dport"))
                self.request_raw_data += HTTP.original
                if self.verify_data('request', self.request_raw_data):
                    self.flag_request = False
                    self.flag_response = True
                    HTTP.original = self.request_raw_data
                    if session_key not in self.data:
                        self.data[session_key] = {'request': [], 'response': []}
                    self.data[session_key]['request'].append(HTTP)
                    self.init_data()
                    return
                else:
                    self.flag_request = True
                    self.flag_response = False
                    self._request = HTTP
                    return

            elif isinstance(HTTP, HTTPResponse):
                try:
                    log.d('HTTPResponse recived: {}'.format(HTTP.original.decode()))
                except:
                    log.d('HTTPResponse recived')
                session_key = "{}-{}-{}-{}".format(IP.fields.get("dst"), IP.fields.get('src'), TCP.fields.get("dport"),
                                                   TCP.fields.get("sport"))

                self.response_raw_data += HTTP.original

                if self.verify_data('response', self.response_raw_data):
                    self.flag_request = False
                    self.flag_response = True
                    if session_key not in self.data:
                        self.data[session_key] = {'request': [], 'response': []}
                    self.data[session_key]['response'].append(HTTP)
                    self.init_data()
                else:
                    self._response = HTTP
                    self.flag_response = True
                    self.flag_request = False
                    return

            if self.flag_request:
                session_key = "{}-{}-{}-{}".format(IP.fields.get("src"), IP.fields.get('dst'), TCP.fields.get("sport"),
                                                   TCP.fields.get("dport"))
                if isinstance(HTTP, Raw):
                    self.request_raw_data += HTTP.original
                try:
                    if self.verify_data('request', self.request_raw_data):
                        self._request.original = self.request_raw_data
                        self.data[session_key]['request'].append(self._request)
                        self.init_data()
                        return
                except dpkt.dpkt.UnpackError:
                    raise
                    # self.init_data()
            elif self.flag_response:
                session_key = "{}-{}-{}-{}".format(IP.fields.get("dst"), IP.fields.get('src'), TCP.fields.get("dport"),
                                                   TCP.fields.get("sport"))
                if isinstance(HTTP, Raw):
                    self.response_raw_data += HTTP.original
                try:
                    if self.verify_data('response', self.response_raw_data):
                        self._response.original = self.response_raw_data
                        self.data[session_key]['response'].append(self._response)
                        self.init_data()
                        return
                    else:
                        return
                except dpkt.dpkt.UnpackError:
                    raise
            # self.init_data()
            # log.d(self.data)


if __name__ == '__main__':
    Capture().run()
